import frappe
from frappe.database import sequence
from frappe import scrub

def patch_sequence():
    """应用达梦数据库序列补丁"""
    
    def dm_get_next_val(doctype_name: str, slug: str = "_id_seq") -> int:
        """达梦数据库获取序列下一个值的实现"""
        sequence_name = scrub(f"{doctype_name}{slug}")
        
        if frappe.db.db_type == "postgres":
            sequence_name = f'"\"{sequence_name}\""'
            sql = f"SELECT nextval({sequence_name})"
        elif frappe.db.db_type == "mariadb":
            sequence_name = f"`{sequence_name}`"
            sql = f"SELECT nextval({sequence_name})"
        elif frappe.db.db_type == "dm":
            sequence_name = f'"{sequence_name}"'
            sql = f"SELECT {sequence_name}.NEXTVAL FROM DUAL"
        else:
            # 默认使用PostgreSQL语法
            sequence_name = f'"\"{sequence_name}\""'
            sql = f"SELECT nextval({sequence_name})"

        try:
            return frappe.db.sql(sql)[0][0]
        except IndexError:
            raise frappe.db.SequenceGeneratorLimitExceeded
    
    def dm_set_next_val(doctype_name: str, slug: str, next_val: int):
        """达梦数据库设置序列下一个值的实现"""
        sequence_name = scrub(f"{doctype_name}{slug}")
        
        if frappe.db.db_type == "dm":
            sql = f'ALTER SEQUENCE "{sequence_name}" RESTART WITH {next_val}'
        else:
            # 调用原始实现
            return sequence._original_set_next_val(doctype_name, slug, next_val)
            
        frappe.db.sql(sql)

    # 保存原始函数
    if not hasattr(sequence, '_original_get_next_val'):
        sequence._original_get_next_val = sequence.get_next_val
    if not hasattr(sequence, '_original_create_sequence'):
        sequence._original_create_sequence = sequence.create_sequence
    if not hasattr(sequence, '_original_set_next_val'):
        sequence._original_set_next_val = sequence.set_next_val
        
    # 应用补丁
    sequence.get_next_val = dm_get_next_val
    sequence.set_next_val = dm_set_next_val

def revert_sequence_patches():
    """恢复序列相关的原始函数"""
    if hasattr(sequence, '_original_get_next_val'):
        sequence.get_next_val = sequence._original_get_next_val
        del sequence._original_get_next_val
        
    if hasattr(sequence, '_original_create_sequence'):
        sequence.create_sequence = sequence._original_create_sequence
        del sequence._original_create_sequence
        
    if hasattr(sequence, '_original_set_next_val'):
        sequence.set_next_val = sequence._original_set_next_val
        del sequence._original_set_next_val
